package net.aiion.weave.impl.resources.url;

import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Iterator;

import net.aiion.weave.spi.resources.ResourceException;
import net.aiion.weave.spi.resources.url.URLConverter;
import net.aiion.weave.spi.resources.url.URLProvider;

public class ClassLoaderURLProvider implements URLProvider
{
    private ClassLoader classLoader;
    private URLConverter urlConverter;

    public ClassLoaderURLProvider() {
        this(Thread.currentThread().getContextClassLoader());
    }

    public ClassLoaderURLProvider(final ClassLoader aClassLoader) {
        classLoader = aClassLoader;
    }

    public void setURLConverter(final URLConverter aURLConverter) {
        urlConverter = aURLConverter;
    }

    public URLConverter getURLConverter() {
        return urlConverter;
    }

    @Override
    public Iterator<URL> iterator() {
        return loadClassLoaderURLs().iterator();
    }

    private Collection<URL> loadClassLoaderURLs() {
        final Collection<URL> classLoaderURLs = new HashSet<>();

        ClassLoader cl = classLoader;
        while (cl != null) {
            classLoaderURLs.addAll(loadClassLoaderURLs(cl));
            cl = cl.getParent();
        }

        return Collections.unmodifiableCollection(classLoaderURLs);
    }

    private Collection<URL> loadClassLoaderURLs(final ClassLoader aClassLoader) {
        if (aClassLoader instanceof URLClassLoader) {
            return loadClassLoaderURLsFromURLClassLoader((URLClassLoader) aClassLoader);
        }

        return loadClassLoaderURLsFromClassLoader(aClassLoader);
    }

    private Collection<URL> loadClassLoaderURLsFromURLClassLoader(final URLClassLoader aURLClassLoader) {
        try {
            return convertURLs(Arrays.asList(aURLClassLoader.getURLs()));
        } catch (final MalformedURLException e) {
            throw new ResourceException(e);
        }
    }

    private Collection<URL> loadClassLoaderURLsFromClassLoader(final ClassLoader aClassLoader) {
        try {
            return loadClassLoaderURLsFromClassLoader1(aClassLoader);
        } catch (final IOException e) {
            throw new ResourceException(e);
        }
    }

    private Collection<URL> loadClassLoaderURLsFromClassLoader1(final ClassLoader aClassLoader) throws IOException {
        final Collection<URL> urls = new ArrayList<>();
        final Enumeration<URL> e = aClassLoader.getResources("");
        while (e.hasMoreElements()) {
            urls.add(convertURL(e.nextElement()));
        }
        return urls;
    }

    private Collection<URL> convertURLs(final Collection<URL> aURLs) throws MalformedURLException {
        if (urlConverter == null) {
            return aURLs;
        }

        final Collection<URL> convertedURLs = new ArrayList<>(aURLs.size());
        for (URL url : aURLs) {
            convertedURLs.add(convertURL(url));
        }
        return convertedURLs;
    }

    private URL convertURL(final URL aURL) throws MalformedURLException {
        if (urlConverter == null) {
            return aURL;
        }

        return urlConverter.convertURL(aURL);
    }
}
